# intermediate_code.py
# 扩展支持 for, do-while, switch, break, continue
from collections import namedtuple
from itertools import count
from parser import run_parser
from lexer import run_lexer

Quad = namedtuple('Quad', ['op','arg1','arg2','res','param_type'])

class TempVarGenerator:
    def __init__(self):
        self._counter = count(1)
    def new(self):
        return f"t{next(self._counter)}"

class QuadGenerator:
    def __init__(self, symtab):
        self.quads = []
        self.temp_gen = TempVarGenerator()
        self.const_map = {}
        self.symtab = symtab   # 引入符号表
        self.loop_stack = []
        self.switch_end = []
        # 当进入函数体或复合语句时切换为 "local"
        self.current_scope = "global"

    def gen(self, op, a1=None, a2=None, res=None, param_type=None):
        # 生成带 param_type 的四元式
        self.quads.append( Quad(op, a1, a2, res, param_type) )

    def _get_string_const(self, s):
        if s not in self.const_map:
            self.const_map[s] = self.temp_gen.new()
        return self.const_map[s]

    def generate(self, ast):
        # 处理声明和函数
        for decl in ast.children:
            if decl.type == 'VarDecl':
                self._gen_global_var_decl(decl)
            elif decl.type=='FunDecl':
                self._gen_fun(decl)
        return self.quads

    def _gen_global_var_decl(self, node):
        """处理全局变量声明"""
        base_type = node.children[0].value  # e.g. 'int'
        for init_decl in node.children[1:]:
            var_name = init_decl.children[0].value

            # 1) 生成 VAR_DECL 四元式
            self.gen('VAR_DECL', var_name, base_type, 'global')

            # 2) 如果有初始化表达式，再生成 ASSIGN
            if len(init_decl.children) == 2:
                rhs = self._gen_expr(init_decl.children[1])
                self.gen('ASSIGN', rhs, None, var_name)

    def _gen_fun(self, node):
        # node.children: [ TypeNode, IDNode, ParamsNode, CompoundStmtNode ]
        name = node.children[1].value
        self.gen('FUNC_BEGIN', name, None, None)

        # --- 新增：对形参生成 VAR_DECL(res='param') ---
        params_node = node.children[2]
        # params_node.type 可能是 'Params'（空），也可能 'ParamList'
        if hasattr(params_node, 'children'):
            for p in params_node.children:
                # 每个 p 是 Param 节点，p.children = [TypeNode, IDNode]
                p_typ = p.children[0].value
                p_name = p.children[1].value
                self.gen('VAR_DECL', p_name, p_typ, 'param')

        # 继续生成函数体
        body = node.children[3]
        self._gen_compound(body)
        self.gen('FUNC_END', name, None, None)

    def _gen_compound(self, node):
        # 进入局部作用域
        prev_scope = self.current_scope
        self.current_scope = "local"

        # 1. 处理局部声明：先产生 VAR_DECL
        for decl in node.children[0].children:  # LocalDecls
            base_type = decl.children[0].value
            for init in decl.children[1:]:
                var_name = init.children[0].value

                # 生成 VAR_DECL 四元式
                self.gen('VAR_DECL', var_name, base_type, 'local')

                # 如有初始化表达式，生成 ASSIGN
                if len(init.children) == 2:
                    rhs = self._gen_expr(init.children[1])
                    self.gen('ASSIGN', rhs, None, var_name)

        # 2. 处理语句列表
        for st in node.children[1].children:
            self._gen_stmt(st)

        # 恢复作用域
        self.current_scope = prev_scope

    def _gen_local_vardecl(self, node):
        """
        处理函数体中混合出现的局部变量声明，
        和 _gen_compound 开头的 LocalDecls 完全一样。
        """
        base_type = node.children[0].value  # e.g. 'int'
        for init in node.children[1:]:
            var_name = init.children[0].value
            # 1) 生成 VAR_DECL
            self.gen('VAR_DECL', var_name, base_type, 'local')
            # 2) 如果有初始化，再生成 ASSIGN
            if len(init.children) == 2:
                rhs = self._gen_expr(init.children[1])
                self.gen('ASSIGN', rhs, None, var_name)

    def _gen_stmt(self,node):
        dispatch = {
            'VarDecl':   self._gen_local_vardecl,
            'ExprStmt': self._gen_exprstmt,
            'IfStmt': self._gen_if,
            'WhileStmt': self._gen_while,
            'DoWhileStmt': self._gen_do_while,
            'ForStmt': self._gen_for,
            'BreakStmt': self._gen_break,
            'ContinueStmt': self._gen_continue,
            'SwitchStmt': self._gen_switch,
            'ReturnStmt': self._gen_return,
            'CompoundStmt': self._gen_compound,
        }
        fn = dispatch.get(node.type)
        if fn: fn(node)

    def _gen_exprstmt(self,node):
        if node.children:
            self._gen_expr(node.children[0])

    def _gen_if(self,node):
        cond, then = node.children[0], node.children[1]
        els = node.children[2] if len(node.children)==3 else None
        Lelse = f"L{self.temp_gen.new()}"
        Lend  = f"L{self.temp_gen.new()}"
        r = self._gen_expr(cond)
        self.gen('JZ',r,None,Lelse)
        self._gen_stmt(then)
        self.gen('JMP',None,None,Lend)
        self.gen('LABEL',None,None,Lelse)
        if els: self._gen_stmt(els)
        self.gen('LABEL',None,None,Lend)

    def _gen_while(self,node):
        cond, body = node.children
        Lstart = f"L{self.temp_gen.new()}"
        Lend   = f"L{self.temp_gen.new()}"
        self.loop_stack.append((Lstart,Lend))
        self.gen('LABEL',None,None,Lstart)
        r = self._gen_expr(cond)
        self.gen('JZ',r,None,Lend)
        self._gen_stmt(body)
        self.gen('JMP',None,None,Lstart)
        self.gen('LABEL',None,None,Lend)
        self.loop_stack.pop()

    def _gen_do_while(self,node):
        body, cond = node.children
        Lstart = f"L{self.temp_gen.new()}"
        Lend   = f"L{self.temp_gen.new()}"
        self.loop_stack.append((Lstart,Lend))
        self.gen('LABEL',None,None,Lstart)
        self._gen_stmt(body)
        r = self._gen_expr(cond)
        self.gen('JNZ',r,None,Lstart)
        self.gen('LABEL',None,None,Lend)
        self.loop_stack.pop()

    def _gen_for(self,node):
        init, cond, step, body = node.children
        if init: self._gen_expr(init)
        Lstart = f"L{self.temp_gen.new()}"
        Lend   = f"L{self.temp_gen.new()}"
        self.gen('LABEL',None,None,Lstart)
        if cond:
            r = self._gen_expr(cond)
            self.gen('JZ',r,None,Lend)
        self.loop_stack.append((Lstart,Lend))
        self._gen_stmt(body)
        if step: self._gen_expr(step)
        self.gen('JMP',None,None,Lstart)
        self.gen('LABEL',None,None,Lend)
        self.loop_stack.pop()

    def _gen_break(self,node):
        _, Lend = self.loop_stack[-1]
        self.gen('JMP',None,None,Lend)

    def _gen_continue(self,node):
        Lstart, _ = self.loop_stack[-1]
        self.gen('JMP',None,None,Lstart)

    def _gen_switch(self,node):
        expr = node.children[0]
        cases = node.children[1:]
        Lend = f"L{self.temp_gen.new()}"
        self.switch_end.append(Lend)
        val = self._gen_expr(expr)
        # 生成每个 case
        for c in cases:
            if c.type=='CaseClause':
                k = self._gen_expr(c.children[0])
                Lcase = f"L{self.temp_gen.new()}"
                self.gen('EQ',val,k,Lcase)
                self.gen('JZ',val,None,Lcase)
                self.gen('LABEL',None,None,Lcase)
                for s in c.children[1:]: self._gen_stmt(s)
            else: # DefaultClause
                for s in c.children: self._gen_stmt(s)
        self.gen('LABEL',None,None,Lend)
        self.switch_end.pop()

    def _gen_return(self,node):
        expr = node.children[0]
        r = self._gen_expr(expr)
        self.gen('RET',r,None,None)

    def _gen_expr(self,node):
        t = node.type
        if t=='STRING':
            cname = self._get_string_const(node.value)
            self.gen('LOAD_STR',node.value,None,cname)
            return cname
        if t=='Assign':
            lhs = node.children[0].value
            rhs = self._gen_expr(node.children[1])
            self.gen('ASSIGN',rhs,None,lhs)
            return lhs
        if t=='BinaryOp':
            L = self._gen_expr(node.children[0])
            op = node.children[1].value
            R = self._gen_expr(node.children[2])
            tmp = self.temp_gen.new()
            self.gen(op,L,R,tmp)
            return tmp
        if t=='UnaryOp':
            tmp = self._gen_expr(node.children[0])
            return tmp  # 可扩展
        if t == 'Call':
            name = node.value
            args = node.children[0].children if node.children else []
            # ...
            if name in ('scanf', 'printf'):
                # 只有一个参数
                if len(args) != 1:
                    raise RuntimeError(f"{name} 需要1个参数")
                arg_node = args[0]
                # 首先生成实际的参数值
                arg = self._gen_expr(arg_node)
                # **记录这个参数节点的类型**：
                ptype = arg_node.type  # 'NUM' or 'STRING' or 'ID' 等
                # 生成 CALL 四元式，最后一位放 param_type
                self.gen('CALL', name, None, arg, param_type=ptype)
                return arg
            
            # 处理普通函数调用
            for a in args:
                v = self._gen_expr(a)
                self.gen('ARG', v, None, None)
            tmp = self.temp_gen.new()
            self.gen('CALL', name, len(args), tmp)
            return tmp
        if t=='ID': return node.value
        if t=='NUM': return node.value
        raise RuntimeError(f"Unknown node {node}")

# 现在还要根据符号表来对变量进行类型检查 TODO 避免识别不了未初始化的变量
def run_generate_code(ast,symtab):
    gen = QuadGenerator(symtab)
    return gen.generate(ast)

if __name__=='__main__':
    code='''
    int main(){
      int a=0;
      for(int i=0;i<10;++i){
        a+=i;
      }
    }
    '''
    tokens,_=run_lexer(code)
    ast=run_parser(tokens)
    quads=run_generate_code(ast)
    for q in quads: print(q)
